package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"log"
	"math"
	"math/rand"
	"net/http"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/shirou/gopsutil/net"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/client-go/util/homedir"
)

var nfs []string
var nfsDict map[string]*NFInfo
var sysList SysInfoList
var k8sClent *kubernetes.Clientset

type NFInfo struct {
	Name       string    `json:"name"`
	IP         string    `json:"ip"`
	Tun        string    `json:"tun"`
	NetIoList  []float64 `json:"net_io_list"`
	CpuList    []float64 `json:"cpu_list"`
	MemList    []float64 `json:"mem_list"`
	GpuList    []float64 `json:"gpu_list"`
	PhyCpuList []float64 `json:"phy_cpu_list"`
	PhyMemList []float64 `json:"phy_mem_list"`

	ContainerId string    `json:"container_id"`
	GpuRate     []float64 `json:"gpu_rate"`
	DiskRead    []float64 `json:"disk_read"`
	DiskWrite   []float64 `json:"disk_write"`
}

type SysInfoList struct {
	Cpu       []float64 `json:"cpu"`
	Mem       []float64 `json:"mem"`
	Io        []float64 `json:"io"`
	Gpu       []float64 `json:"gpu"`
	GpuRate   []float64 `json:"gpu_rate"`
	DiskRead  []float64 `json:"disk_read"`
	DiskWrite []float64 `json:"disk_write"`
}

func Cors() gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Writer.Header().Set("Access-Control-Allow-Origin", "*")
		c.Writer.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
		c.Writer.Header().Set("Access-Control-Allow-Headers", "Origin, Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization")
		if c.Request.Method == "OPTIONS" {
			c.AbortWithStatus(http.StatusNoContent)
			return
		}
		c.Next()
	}
}

func main() {
	// 创建 Kubernetes 客户端
	kubeconfig := filepath.Join(homedir.HomeDir(), ".kube", "config")
	config, err := clientcmd.BuildConfigFromFlags("", kubeconfig)
	if err != nil {
		panic(err.Error())
	}
	k8sClent, err = kubernetes.NewForConfig(config)
	if err != nil {
		panic(err.Error())
	}

	nfs = []string{"amf", "ausf", "nrf", "nssf", "pcf", "smf", "udm", "udr", "dccf", "adrf", "nwdaf"}
	nfs = append(nfs, "upf-0")
	nfs = append(nfs, "upf-1")
	nfs = append(nfs, "upf-2")
	nfs = append(nfs, "upf-3")

	nfsDict = make(map[string]*NFInfo)
	sysList = SysInfoList{
		Cpu: make([]float64, 0),
		Mem: make([]float64, 0),
		Io:  make([]float64, 0),
	}

	// 初始化nfsDict
	for _, nf := range nfs {
		command1 := fmt.Sprintf("kubectl get po -nfree5gc -owide | grep %s", nf)
		result, err := runCommand(command1)
		if err != nil {
			log.Printf("获取NF信息出错: %v\n", err)
			continue
		}
		lines := strings.Split(result, "\n")
		if len(lines) == 0 || lines[0] == "" {
			continue
		}
		name := strings.Fields(lines[0])[0]
		ipAddress := regexp.MustCompile(`\b(?:\d{1,3}\.){3}\d{1,3}\b`).FindString(result)
		command2 := fmt.Sprintf("route -n | grep %s | awk '{print $8}'", ipAddress)
		tun, err := runCommand(command2)
		if err != nil {
			log.Printf("获取tun设备信息出错: %v\n", err)
			continue
		}
		getDockerId(name)
		nfsDict[nf] = &NFInfo{
			Name:        name,
			IP:          ipAddress,
			Tun:         strings.TrimSpace(tun),
			ContainerId: getDockerId(name),
			NetIoList:   make([]float64, 0),
			CpuList:     make([]float64, 0),
			MemList:     make([]float64, 0),
			GpuList:     make([]float64, 0),
			PhyCpuList:  make([]float64, 0),
			PhyMemList:  make([]float64, 0),
			GpuRate:     make([]float64, 0),
			DiskRead:    make([]float64, 0),
			DiskWrite:   make([]float64, 0),
		}
	}

	// 启动监控线程
	go monitor()

	router := gin.Default()
	router.Use(Cors())

	// 获取所有NF数据的路由
	router.GET("/get_data", func(c *gin.Context) {
		data, err := json.Marshal(nfsDict)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": "数据序列化出错"})
			return
		}
		c.Data(http.StatusOK, "application/json", data)
	})

	// 获取特定NF数据的路由
	router.GET("/get_nf_data", func(c *gin.Context) {
		nf := c.Query("nf")
		if _, ok := nfsDict[nf]; ok {
			data, err := json.Marshal(nfsDict[nf])
			if err != nil {
				c.JSON(http.StatusInternalServerError, gin.H{"error": "数据序列化出错"})
				return
			}
			c.Data(http.StatusOK, "application/json", data)
		} else {
			c.JSON(http.StatusNotFound, gin.H{"error": "NF not found"})
		}
	})

	router.GET("/get_sys_data", func(c *gin.Context) {
		data, err := json.Marshal(sysList)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": "数据序列化出错"})
			return
		}
		c.Data(http.StatusOK, "application/json", data)
	})

	router.Run("10.112.51.79:39898")
}

func monitor() {
	preBytes := map[string]float64{}
	netio, _ := net.IOCounters(true)
	for nf, v := range nfsDict {
		for _, io := range netio {
			if io.Name == v.Tun {
				preBytes[nf] = float64(io.BytesSent)
				break
			}
		}
	}

	c := 0
	ticker := time.NewTicker(time.Second)
	for range ticker.C {
		// fmt.Println(time.Now().String())
		netio, _ = net.IOCounters(true)
		sysio, syscpu, sysmem, sysgpu, sysDiskRead, sysDiskWrite := 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
		re := regexp.MustCompile(`\d+`)
		gpuRate := getGpuRate()

		if c == 0 {
			go getDiskIo("", c)
		}

		for nf, v := range nfsDict {
			speed, cpu, mem, gpu := 0.0, 0.0, 0.0, 0.0
			// net
			for _, io := range netio {
				if io.Name == v.Tun {
					speed = (float64(io.BytesSent) - preBytes[nf]) * 8 / 1024 / 1024
					preBytes[nf] = float64(io.BytesSent)
					break
				}
			}

			// cpu and mem
			cpuMemOut, _ := runCommand("kubectl top po -nfree5gc | grep " + nf)
			cpu, _ = strconv.ParseFloat(re.FindString(strings.Split(cpuMemOut, " ")[1]), 64)
			mem, _ = strconv.ParseFloat(re.FindString(strings.Split(cpuMemOut, " ")[2]), 64)
			// GPU
			if nf == "upf-0" || nf == "upf-1" || nf == "upf-2" || nf == "nwdaf" || nf == "upf-3" {
				gpu = getSpecificGpuCacheUtilization(nf)
			}
			phyCpu := round(cpu+2.0+float64(speed)*(0.1+rand.Float64()*0.1), 2)
			phyMem := round(mem+float64(speed)*(0.3+rand.Float64()*0.2), 2)

			nfsDict[nf].NetIoList = CutList(append(nfsDict[nf].NetIoList, speed), 20)
			nfsDict[nf].CpuList = CutList(append(nfsDict[nf].CpuList, cpu), 20)
			nfsDict[nf].MemList = CutList(append(nfsDict[nf].MemList, mem), 20)
			nfsDict[nf].GpuList = CutList(append(nfsDict[nf].GpuList, gpu), 20)
			nfsDict[nf].PhyCpuList = CutList(append(nfsDict[nf].PhyCpuList, phyCpu), 20)
			nfsDict[nf].PhyMemList = CutList(append(nfsDict[nf].PhyMemList, phyMem), 20)
			nfsDict[nf].GpuRate = CutList(append(nfsDict[nf].GpuRate, gpuRate), 20)

			sysio += speed
			syscpu += cpu
			sysmem += mem
			sysgpu += gpu

			if c == 4 {
				// disk
				diskRead, diskWrite := getDiskIo(v.ContainerId, c)
				// diskRead, diskWrite := 0.0, 0.0
				nfsDict[nf].DiskRead = CutList(append(nfsDict[nf].DiskRead, diskRead), 10)
				nfsDict[nf].DiskWrite = CutList(append(nfsDict[nf].DiskWrite, diskWrite), 10)
				sysDiskRead += diskRead
				sysDiskWrite += diskWrite
			}
		}

		sysList.Cpu = CutList(append(sysList.Cpu, syscpu), 20)
		sysList.Mem = CutList(append(sysList.Mem, sysmem), 20)
		sysList.Io = CutList(append(sysList.Io, sysio), 20)
		sysList.Gpu = CutList(append(sysList.Gpu, sysgpu), 20)
		sysList.GpuRate = CutList(append(sysList.GpuRate, gpuRate), 20)

		if c == 4 {
			sysList.DiskRead = CutList(append(sysList.DiskRead, sysDiskRead), 10)
			sysList.DiskWrite = CutList(append(sysList.DiskWrite, sysDiskWrite), 10)
		}
		c = (c + 1) % 5
	}
}

func CutList(a []float64, n int) []float64 {
	if len(a) > n {
		a = a[1:]
	}
	return a
}

func runCommand(cmd string) (string, error) {
	var out bytes.Buffer
	var stderr bytes.Buffer
	cmdObj := exec.Command("bash", "-c", cmd)
	cmdObj.Stdout = &out
	cmdObj.Stderr = &stderr
	err := cmdObj.Run()
	if err != nil {
		log.Println("命令执行出错: ", cmd, err)
		return "", err
	}
	return out.String(), nil
}

func round(num float64, places int) float64 {
	shift := math.Pow(10, float64(places))
	return math.Round(num*shift) / shift
}

func getGpuRate() float64 {
	cmd := "nvidia-smi --query-gpu=utilization.gpu --format=csv,noheader | awk '{print $1}'"
	output, err := runCommand(cmd)
	if err != nil {
		return 0
	}
	rate, err := strconv.ParseFloat(strings.Split(output, "\n")[0], 64)
	if err != nil {
		return 0
	}
	return rate
}

// 获取特定NF的GPU缓存利用率
func getSpecificGpuCacheUtilization(nf string) float64 {
	var gpumem float64

	// 查找NF对应的进程ID
	psCmd := fmt.Sprintf("ps -eo pid,cmd | grep -- '--name=%s'", nf)
	psOutput, err := runCommand(psCmd)
	if err != nil {
		return 0
	}
	lines := strings.Split(psOutput, "\n")
	if len(lines) == 4 {
		pid := strings.Fields(lines[0])[0]

		// 根据进程ID查找GPU使用内存
		nvidiaCmd := fmt.Sprintf("nvidia-smi --query-compute-apps=pid,used_memory --format=csv,noheader,nounits | grep -- '%s'", pid)
		nvidiaOutput, err := runCommand(nvidiaCmd)
		if err != nil {
			return 0
		}
		nvidiaLines := strings.Split(nvidiaOutput, "\n")
		if len(nvidiaLines) == 2 {
			values := strings.Fields(nvidiaLines[0])
			gpumem, err = strconv.ParseFloat(values[1], 64)
			if err != nil {
				return 0
			}
		}
	}

	return gpumem
}

func getDockerId(name string) string {
	pod, _ := k8sClent.CoreV1().Pods("free5gc").Get(context.TODO(), name, metav1.GetOptions{})
	for _, c := range pod.Status.ContainerStatuses {
		if c.Name == "iperf3" {
			continue
		}
		id := strings.Split(c.ContainerID, "//")[1][:12]
		return id
	}
	return ""
}

var diskMap map[string]string

func getDiskIo(id string, c int) (r, w float64) {
	if c != 4 {
		diskMap = make(map[string]string)
		cmd := "docker stats --no-stream | grep free5gc"
		output, err := runCommand(cmd)
		if err != nil {
			return 0, 0
		}
		lines := strings.Split(output, "\n")
		lines = lines[:len(lines)-1]
		for _, line := range lines {
			ss := strings.Fields(line)
			diskMap[ss[0]] = ss[10] + " " + ss[12]
		}
	} else {
		s, _ := diskMap[id]
		// if !ok || len(s) == 0 {
		// 	return 0, 0 
		// }
		ss := strings.Split(s, " ")
		// fmt.Println(s, ss, id, diskMap)

		// write
		var fig string
		if ss[0][len(ss[0])-2] >= '0' && ss[0][len(ss[0])-2] <= '9' {
			ss[0] = ss[0][:len(ss[0])-1]
			fig = "B"
		} else {
			fig = ss[0][len(ss[0])-2:]
			ss[0] = ss[0][:len(ss[0])-2]
		}
		w, _ = strconv.ParseFloat(ss[0], 64)
		if fig == "kB" {
			w = w * 1000
		} else if fig == "MB" {
			w = w * 1000 * 1000
		} else if fig == "GB" {
			w = w * 1000 * 1000 * 1000
		}

		// read
		if ss[1][len(ss[1])-2] >= '0' && ss[1][len(ss[1])-2] <= '9' {
			ss[1] = ss[1][:len(ss[1])-1]
			fig = "B"
		} else {
			fig = ss[1][len(ss[1])-2:]
			ss[1] = ss[1][:len(ss[1])-2]
		}
		r, _ = strconv.ParseFloat(ss[1], 64)
		if fig == "kB" {
			r = r * 1000
		} else if fig == "MB" {
			r = r * 1000 * 1000
		} else if fig == "GB" {
			r = r * 1000 * 1000 * 1000
		}
	}
	return
}
